# 1. 定义循环状态：新增“是否继续对话”字段
from typing import TypedDict, List

from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import StateGraph, END


class LoopChatState(TypedDict):
    messages: List[HumanMessage | AIMessage]
    continue_chat: bool  # 标记是否继续对话
    user_next_input: str  # 存储用户下一轮输入


# 2. 定义“询问是否继续”节点
def ask_continue_node(state: LoopChatState) -> LoopChatState:
    """生成回答后，询问用户是否继续对话"""
    # 在AI回答后，新增一条“询问消息”
    continue_msg = AIMessage(content="是否需要继续提问？（输入‘结束’终止对话，其他内容继续）")
    return {
        "messages": state["messages"] + [continue_msg],
        "continue_chat": True,  # 初始默认继续
        "user_next_input": state["user_next_input"]  # 保留原有值
    }


# 3. 定义“获取用户输入”节点（模拟用户输入，实际项目中可对接前端）
def get_user_input_node(state: LoopChatState) -> LoopChatState:
    """获取用户下一轮输入，判断是否继续"""
    user_input = input("请输入你的回复：")  # 模拟前端输入
    # 判断是否终止：若输入“结束”，则设置 continue_chat=False
    continue_chat = user_input != "结束"
    # 更新状态：新增用户输入，更新 continue_chat
    return {
        "messages": state["messages"] + [HumanMessage(content=user_input)],
        "continue_chat": continue_chat,
        "user_next_input": user_input
    }


# 4. 初始化图
graph = StateGraph(LoopChatState)

from langchain.chat_models import init_chat_model

# 初始化LLM（请确保已安装对应依赖并配置API密钥）
llm = init_chat_model(model="deepseek-chat", model_provider="deepseek")


def llm_answer_node(state: LoopChatState) -> LoopChatState:
    """LLM回答节点：基于对话历史生成回答"""
    messages = state["messages"]
    # 调用LLM生成回答
    response = llm.invoke(messages)
    # 更新消息列表
    return {
        "messages": messages + [response],
        "continue_chat": state["continue_chat"],
        "user_next_input": state["user_next_input"]
    }


# 5. 添加节点
graph.add_node("llm_answer", llm_answer_node)  # LLM回答节点
graph.add_node("ask_continue", ask_continue_node)  # 询问是否继续节点
graph.add_node("get_user_input", get_user_input_node)  # 获取用户输入节点

# 6. 定义边：实现循环逻辑
# 入口 → LLM节点（第一轮先生成回答）
graph.set_entry_point("llm_answer")

# LLM节点 → 询问是否继续（生成回答后，询问用户）
graph.add_edge("llm_answer", "ask_continue")

# 询问是否继续 → 获取用户输入（无条件）
graph.add_edge("ask_continue", "get_user_input")


# 条件判断函数：根据状态决定下一步流向
def decide_next_step(state: LoopChatState) -> str:
    if state["continue_chat"]:
        return "llm_answer"  # 继续对话，回到LLM节点
    else:
        return END  # 终止对话，进入出口


# 获取用户输入 → 分支判断（通过条件函数路由）
graph.add_conditional_edges(
    "get_user_input",
    decide_next_step  # 传入条件判断函数
)

# 7. 编译并运行循环
compiled_graph = graph.compile()
initial_state = {
    "messages": [HumanMessage(content="LangGraph 支持循环吗？")],
    "continue_chat": True,
    "user_next_input": ""
}

# 执行（会持续等待用户输入，直到输入“结束”）
result = compiled_graph.invoke(initial_state)

# 打印最终结果
print("\n对话结束，最终消息历史：")
for msg in result["messages"]:
    print(f"{msg.type}: {msg.content}")
